#coding=utf-8
import pymongo
import jieba.posseg as pseg


class Snippet(object):

    def connection(self, database):
        """
        数据库连接
        :param name: 数据库名
        :return: 数据库的collection
        """
        connection = pymongo.MongoClient('localhost', 27017)
        db = connection['news']
        collection = db[database]
        return collection

    def read_news(self, database, news_id):
        """
        由新闻id从数据库中获取到新闻内容
        :param news_id: 新闻id号
        :return: 新闻内容
        """
        collection = self.connection(database)
        news = list(collection.find({'news_id': news_id}))[0]
        return news

    def create_snippet(self, database, query, news_id, window=100, query_list=None):
        """
        snippet生成
        :param database: 数据库
        :param query: 用户输入的查询
        :param news_id: 该查询对应的文档id号
        :param window: snippet 窗口大小
        :param query_list: 用户输入的查询转换成list（可省）
        :return: 生成的snippet
        """
        # 如果包含完整的query
        flag = False
        news = self.read_news(database, news_id)
        time = news['news_pub_time']
        content = news['news_content']
        index = 1
        for i in range(1, len(content)):
            if query in content[i]:
                index = i
                flag = True
                break

        # 如果不包含完整的query
        if not flag:
            if query_list is []:  # 如果query_list为空，则自己分词
                stopwords = ['uj', 'ul', 'x', 'eng']
                words = pseg.cut(query)
                for word, tag in words:
                    if tag not in stopwords:
                        query_list.append(word)

            # 获取包含最多关键词的句子作为snippet的首句
            max_count = 0
            for i in range(1, len(content)):
                count = 0
                for query_word in query_list:
                    if query_word in content[i]:
                        count += 1
                if count > max_count:
                    index = i

        # snippet 生成
        news_snippet = content[index]
        while len(news_snippet) < window and index + 1 < len(content) - 1:
            index += 1
            news_snippet += content[index]
        news_snippet = time + ' - ' + news_snippet[0:window] + '...'
        return news_snippet

    def create_preview(self, database, news_id):
        news = self.read_news(database, news_id)
        return news['news_title'], ' '.join(news['news_content'][1:])


if __name__ == '__main__':
    snippet = Snippet()
    # create_snippet(self, database, query, news_id, window, query_list=None)
    # 生成snippet
    news_snippet = snippet.create_snippet('wangyi_news', '国务院任免', 'D55N97190001875N')
    print(news_snippet)
    
    # 生成预览
    title, preview = snippet.create_preview('wangyi_news', 'D55N97190001875N')
    print(title, preview)


